/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "velox/expression/ExprRewriteRegistry.h"
#include "velox/expression/FunctionSignature.h"

namespace facebook::velox::expression {

void ExprRewriteRegistry::registerRewrite(ExpressionRewrite rewrite) {
  registry_.withWLock([&](auto& list) { list.push_back(std::move(rewrite)); });
}

void ExprRewriteRegistry::clear() {
  registry_.withWLock([&](auto& list) { list.clear(); });
}

core::TypedExprPtr ExprRewriteRegistry::rewrite(
    const core::TypedExprPtr& expr) {
  const auto& inputs = expr->inputs();
  VELOX_CHECK(
      std::any_of(
          inputs.begin(),
          inputs.end(),
          [](const core::TypedExprPtr& input) {
            return !input->isConstantKind();
          }),
      "Expression should have at least one non-constant input.");

  core::TypedExprPtr result = expr;
  registry_.withRLock([&](const auto& list) {
    for (const auto& rewrite : list) {
      VELOX_CHECK_NOT_NULL(rewrite);
      if (auto rewritten = (rewrite)(expr)) {
        result = rewritten;
        break;
      }
    }
  });

  return result;
}
} // namespace facebook::velox::expression
